from flask import Flask, request, jsonify
import os
import json
from generate_prompt_js import parse_jsonl, get_prompt
from text_embedding import TextBase
import zhipu_question2response as zhipu_qr
import deepseek_question2response as deepseek_qr
from post_extract import post_extract_js

# 初始化 Flask 应用
app = Flask(__name__)

# 全局变量和配置
level_data_file = "lulu_exp/sample_levels.jsonl"
base_save_file = "lulu_exp/id_text_vector.parquet"
save_file = "lulu_exp/output_log.txt"
deep_seek_name = "DeepSeek"

# 预加载数据
level_datas = parse_jsonl(level_data_file)

# 加载新增数据
import glob

for file in glob.glob("lulu_exp/new_levels/*.json"):
    with open(file, "r", encoding="utf-8") as f:
        data = json.load(f)
        level_datas.append({
            "query_str": data["desprition"],
            "js_code": data["code"],
            "title": data["title"]
        })

text_base = TextBase()
query_texts = [d["query_str"] for d in level_datas]

# 构建或加载基础数据
if os.path.exists(base_save_file):
    text_base.load(base_save_file)
    rebuild = any(text_base.strong_match(query_text) is None for query_text in query_texts)
else:
    rebuild = True

if rebuild:
    text_base = TextBase()
    text_base.build_base(query_texts)
    text_base.save(base_save_file)


# 工具函数
def generate_prompt(query):
    """生成 Prompt"""
    prompt = get_prompt(query, text_base, level_datas, None, if_english=False)
    return prompt

def get_response(prompt, model_choice):
    """根据模型生成 Response"""
    if model_choice == "Zhipu":
        response = zhipu_qr.question2response(prompt)
    elif model_choice == deep_seek_name:
        response = deepseek_qr.question2response(prompt)
    else:
        raise ValueError("Invalid model choice")
    
    # 保存日志
    with open(save_file, "a", encoding="utf-8") as f:
        f.write(json.dumps({"prompt": prompt, "response": response}, ensure_ascii=False) + "\n")
    
    return response

def extract_js_code(response):
    """从 Response 中提取 JS Code"""
    data = post_extract_js(response)
    return data.get("js_code", "")


# Flask 路由
@app.route('/generate_code', methods=['POST'])
def generate_code():
    try:
        # 获取请求数据
        body = request.json
        query = body.get('query')
        model_choice = body.get('model')

        if not query or not model_choice:
            return jsonify({"error": "Missing 'query' or 'model' field"}), 400

        # 生成 Prompt
        prompt = generate_prompt(query)

        # 获取 Response
        response = get_response(prompt, model_choice)

        # 提取 JS Code
        js_code = extract_js_code(response)

        # 返回结果
        return jsonify({
            "prompt": prompt,
            "response": response,
            "js_code": js_code
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500


# 启动 Flask 应用
if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000, debug=True)
